# -*- coding: utf-8 -*-
"""
GeoPackage文件处理器
将GeoPackage文件转换为标准化的PostGIS导入参数格式
"""

import geopandas as gpd
import pandas as pd
import json
from typing import Dict, Any, Optional, List
import os
from loguru import logger
from shapely.geometry import MultiPolygon, Point


class GeoPackageProcessor:
    """GeoPackage文件处理器，转换为PostGIS导入格式"""
    
    def __init__(self, file_path: str):
        self.file_path = file_path
        self.gdf = None
        self.layer_name = None
        
    def list_layers(self) -> List[str]:
        """列出GeoPackage中的所有图层"""
        try:
            import fiona
            layers = fiona.listlayers(self.file_path)
            return layers
        except Exception as e:
            logger.error(f"无法列出图层: {str(e)}")
            return []
    
    def load_data(self, layer_name: str = None) -> Dict[str, Any]:
        """
        加载GeoPackage文件并转换为标准格式
        
        Args:
            layer_name: 图层名称，如果为None则使用第一个图层
            
        Returns:
            Dict包含: {
                'data': 标准化后的GeoDataFrame,
                'geometry_type': 几何类型,
                'crs': 坐标系,
                'columns': 字段列表,
                'record_count': 记录数量,
                'feature_count': 原始特征数量,
                'layer_name': 实际使用的图层名称,
                'available_layers': 可用图层列表
            }
        """
        try:
            # 获取可用图层
            available_layers = self.list_layers()
            if not available_layers:
                raise ValueError("GeoPackage中没有找到可用图层")
                
            # 确定图层名称
            if layer_name is None:
                self.layer_name = available_layers[0]
            elif layer_name not in available_layers:
                raise ValueError(f"图层 '{layer_name}' 不存在。可用图层: {available_layers}")
            else:
                self.layer_name = layer_name
                
            # 读取指定图层
            self.gdf = gpd.read_file(self.file_path, layer=self.layer_name)
            
            # 标准化数据结构
            standardized_data = self._standardize_data()
            
            result = {
                'data': standardized_data,
                'geometry_type': self._get_geometry_type(),
                'crs': str(self.gdf.crs) if self.gdf.crs else 'EPSG:4326',
                'columns': list(standardized_data.columns),
                'record_count': len(standardized_data),
                'feature_count': len(self.gdf),
                'layer_name': self.layer_name,
                'available_layers': available_layers,
                'source_format': 'GeoPackage'
            }
            
            logger.info(f"GeoPackage处理完成: {os.path.basename(self.file_path)}")
            logger.info(f"图层: {self.layer_name}, 特征数: {result['feature_count']}, 几何类型: {result['geometry_type']}")
            
            return result
            
        except Exception as e:
            logger.error(f"GeoPackage文件处理失败: {e}")
            raise
    
    def _standardize_data(self) -> gpd.GeoDataFrame:
        """标准化数据结构"""
        if self.gdf is None:
            raise ValueError("请先加载数据")
            
        # 创建标准化列
        standardized = self.gdf.copy()
        
        # GeoPackage文件通常不包含嵌套的properties字段，这是GeoJSON特有的
        
        # 确保有标准字段
        if 'code' not in standardized.columns:
            # 优先尝试stcd字段作为code，然后尝试其他常见水文字段
            code_candidates = ['stcd', 'STCD', 'stcd_field', 'gd', 'GD', 'id', 'ID', 'code', 'Code', 'CODE', 'station_id', 'StationID']
            for candidate in code_candidates:
                if candidate in standardized.columns:
                    standardized['code'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为code")
                    break
            else:
                standardized['code'] = [f"GPKG_{i:06d}" for i in range(len(standardized))]
                logger.warning("未找到合适的code字段，使用自动生成")

        if 'name' not in standardized.columns:
            # 尝试从更多可能的名称字段获取
            name_candidates = ['name', 'Name', 'NAME', 'stnm', 'STNM', 'title', 'label', 'station_name', 'reservoir_name', 'lake_name']
            for candidate in name_candidates:
                if candidate in standardized.columns:
                    standardized['name'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为name")
                    break
            else:
                # 如果没有合适的名称字段，尝试使用code作为name
                if 'code' in standardized.columns:
                    standardized['name'] = standardized['code']
                    logger.info("使用code作为name")
                else:
                    standardized['name'] = [f"未命名_{i}" for i in range(len(standardized))]
                
        if 'type' not in standardized.columns:
            # 尝试从类型相关字段获取
            type_candidates = ['type', 'Type', 'TYPE', 'category', 'Category', 'kind', 'Kind', 'reservoir_type', 'station_type']
            for candidate in type_candidates:
                if candidate in standardized.columns:
                    standardized['type'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为type")
                    break
            else:
                # 根据几何类型推断合理的默认值
                first_geom_type = standardized.geometry.iloc[0].geom_type if len(standardized) > 0 else 'Unknown'
                if first_geom_type in ['Point', 'MultiPoint']:
                    standardized['type'] = '水文站点'
                elif first_geom_type in ['Polygon', 'MultiPolygon']:
                    standardized['type'] = '水库'
                else:
                    standardized['type'] = '要素'
        else:
            # 处理空值、null值和空字符串
            mask = standardized['type'].isna() | (standardized['type'] == '') | (standardized['type'].astype(str).str.strip() == '')
            standardized.loc[mask, 'type'] = '要素'
            
        # 多余字段直接丢弃，extra_data保持为空
        standardized['extra_data'] = ''
        
        # 转换几何类型以匹配数据库要求
        standardized = self._convert_geometry_types(standardized)
        
        # 转换坐标系为EPSG:4326
        if standardized.crs is None:
            standardized = standardized.set_crs('EPSG:4326')
        elif str(standardized.crs) != 'EPSG:4326':
            standardized = standardized.to_crs('EPSG:4326')
            
        # 只保留标准字段
        result = standardized[['code', 'name', 'type', 'extra_data', 'geometry']].copy()
        
        return result
    
    def _convert_geometry_types(self, gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
        """转换几何类型以匹配数据库要求"""
        if gdf is None or len(gdf) == 0:
            return gdf
            
        # 创建副本避免修改原始数据
        result = gdf.copy()
        
        # 获取原始几何类型
        first_geom_type = result.geometry.iloc[0].geom_type
        
        # 根据类型进行转换
        if first_geom_type == 'Polygon':
            # 将Polygon转换为MultiPolygon
            result.geometry = result.geometry.apply(
                lambda geom: geom if geom is None else 
                geom if geom.geom_type == 'MultiPolygon' else
                geom if geom.geom_type != 'Polygon' else
                MultiPolygon([geom])
            )
        elif first_geom_type == 'MultiPoint':
            # 将MultiPoint转换为Point（取第一个点）
            result.geometry = result.geometry.apply(
                lambda geom: geom if geom is None else
                geom if geom.geom_type == 'Point' else
                geom if geom.geom_type != 'MultiPoint' else
                list(geom.geoms)[0] if len(list(geom.geoms)) > 0 else None
            )
            
        return result
    
    def _get_geometry_type(self) -> str:
        """获取几何类型"""
        if self.gdf is None or len(self.gdf) == 0:
            return 'Unknown'
            
        geom_type = self.gdf.geometry.iloc[0].geom_type
        
        # 标准化几何类型
        type_mapping = {
            'Polygon': 'MULTIPOLYGON',
            'MultiPolygon': 'MULTIPOLYGON',
            'Point': 'POINT',
            'MultiPoint': 'POINT'
        }
        
        return type_mapping.get(geom_type, geom_type.upper())
    
    def validate_data(self) -> Dict[str, Any]:
        """验证数据完整性"""
        if self.gdf is None:
            return {'valid': False, 'errors': ['数据未加载']}
            
        errors = []
        warnings = []
        
        # 检查几何有效性
        invalid_geoms = ~self.gdf.geometry.is_valid
        if invalid_geoms.any():
            errors.append(f"发现{invalid_geoms.sum()}个无效几何图形")
            
        # 检查空几何
        empty_geoms = self.gdf.geometry.is_empty
        if empty_geoms.any():
            errors.append(f"发现{empty_geoms.sum()}个空几何图形")
            
        # 检查重复code
        if 'code' in self.gdf.columns:
            duplicate_codes = self.gdf['code'].duplicated().sum()
            if duplicate_codes > 0:
                errors.append(f"发现{duplicate_codes}个重复code")
                
        # 检查坐标系
        if self.gdf.crs is None:
            warnings.append("未定义坐标系，将使用EPSG:4326")
            
        # 检查数据类型
        for col in self.gdf.columns:
            if col != 'geometry':
                dtype = str(self.gdf[col].dtype)
                if 'object' in dtype and self.gdf[col].notna().any():
                    # 检查是否有复杂对象类型
                    sample = self.gdf[col].dropna().iloc[0]
                    if isinstance(sample, (dict, list)):
                        warnings.append(f"字段 '{col}' 包含复杂对象类型，将被序列化为JSON")
        
        return {
            'valid': len(errors) == 0,
            'errors': errors,
            'warnings': warnings
        }
    
    def get_layer_info(self) -> Dict[str, Any]:
        """获取图层详细信息"""
        if self.gdf is None:
            return {}
            
        return {
            'bbox': self.gdf.total_bounds.tolist(),
            'crs': str(self.gdf.crs),
            'geometry_types': self.gdf.geometry.geom_type.value_counts().to_dict(),
            'field_types': self.gdf.dtypes.to_dict(),
            'memory_usage': self.gdf.memory_usage(deep=True).sum(),
            'null_counts': self.gdf.isnull().sum().to_dict()
        }

    def get_metadata(self) -> Dict[str, Any]:
        """获取文件元数据"""
        return self.get_layer_info()
    
    def preview_layers(self) -> Dict[str, Dict[str, Any]]:
        """预览所有图层的基本信息"""
        layers = self.list_layers()
        preview = {}
        
        for layer in layers:
            try:
                temp_gdf = gpd.read_file(self.file_path, layer=layer, limit_rows=5)
                preview[layer] = {
                    'geometry_type': temp_gdf.geometry.iloc[0].geom_type if len(temp_gdf) > 0 else 'Unknown',
                    'record_count': len(gpd.read_file(self.file_path, layer=layer, limit_rows=0)),
                    'columns': list(temp_gdf.columns),
                    'sample_count': len(temp_gdf)
                }
            except Exception as e:
                preview[layer] = {'error': str(e)}
                
        return preview


# 测试示例
if __name__ == "__main__":
    # 创建测试GeoPackage数据
    test_data = gpd.GeoDataFrame({
        'name': ['测试水库1', '测试水库2'],
        'capacity': [1000, 500],
        'type': ['大型', '中型'],
        'code': ['RSV001', 'RSV002']
    })
    
    # 添加几何数据
    from shapely.geometry import Point
    test_data = test_data.set_geometry([
        Point(116.4074, 39.9042),
        Point(117.1993, 39.0851)
    ])
    
    # 保存测试文件
    test_file = "test_data.gpkg"
    test_data.to_file(test_file, driver="GPKG", layer="reservoirs")
    
    try:
        processor = GeoPackageProcessor(test_file)
        
        # 列出所有图层
        layers = processor.list_layers()
        print("=== 可用图层 ===")
        for layer in layers:
            print(f"- {layer}")
            
        # 预览图层
        preview = processor.preview_layers()
        print("\n=== 图层预览 ===")
        for layer, info in preview.items():
            print(f"{layer}: {info}")
        
        # 加载并处理数据
        if layers:
            result = processor.load_data(layers[0])
            
            print("\n--- load_data() 返回的字典 ---")
            print(result)
            print("--- 字典结束 ---")

            # 验证数据
            validation = processor.validate_data()
            
            print("\n=== GeoPackage处理结果 ===")
            print(f"文件: {test_file}")
            print(f"图层: {result['layer_name']}")
            print(f"记录数: {result['record_count']}")
            print(f"几何类型: {result['geometry_type']}")
            print(f"坐标系: {result['crs']}")
            print(f"字段: {result['columns']}")
            print(f"数据验证: {'通过' if validation['valid'] else '失败'}")
            
            if validation['warnings']:
                print("警告:", validation['warnings'])
            if not validation['valid']:
                print("错误:", validation['errors'])
                
            # 显示元数据
            metadata = processor.get_layer_info()
            print("\n=== 图层元数据 ===")
            print(f"边界框: {metadata.get('bbox', 'N/A')}")
            print(f"几何类型分布: {metadata.get('geometry_types', {})}")
            
            # 显示数据示例
            print("\n=== 数据示例 ===")
            print(result['data'][['code', 'name', 'type', 'extra_data']].head())
        
    finally:
        # 清理测试文件
        if os.path.exists(test_file):
            os.remove(test_file)

    print("\n" + "="*50)
    print("=== 本地GeoPackage文件测试 ===")
    print("="*50)
    
    # 本地测试：直接测试gansu.gpkg
    test_file = "gansu.gpkg"
    if os.path.exists(test_file):
        try:
            print(f"\n正在测试本地文件: {test_file}")
            processor = GeoPackageProcessor(test_file)
            
            # 列出图层
            layers = processor.list_layers()
            print(f"发现图层: {layers}")
            
            if layers:
                # 加载第一个图层
                result = processor.load_data(layers[0])

                print("\n--- load_data() 返回的字典 ---")
                print(result)
                print("--- 字典结束 ---\n")

                validation = processor.validate_data()
                
                print(f"✅ 成功处理: {result['feature_count']} 条记录")
                print(f"几何类型: {result['geometry_type']}")
                print(f"坐标系: {result['crs']}")
                print(f"数据验证: {'通过' if validation['valid'] else '失败'}")
                
                if validation['warnings']:
                    print(f"⚠️ 警告: {validation['warnings']}")
                
                # 显示前3条数据
                print("数据样本:")
                print(result['data'][['code', 'name']].head(3).to_string(index=False))
                
        except Exception as e:
            print(f"❌ 处理失败: {str(e)}")
    else:
        print(f"❌ 未找到文件: {test_file}")